import numpy as np
import cv2 as cv
import math
import time


def depth_inpaint(img):
    """inpaint图像修复"""
    print('inpaint processing...')
    t1 = time.time()
    mask = (img==0).astype(np.uint8)
    dst = cv.inpaint(img, mask, 3, cv.INPAINT_TELEA)
    time_used = (time.time() - t1)*1000
    print('Use {:.2f} ms.'.format(time_used))
    return dst


def depth_medianblur(img, kernel=2):
    """中值滤波
    :param img:
    :param kernel:核大小,1为3x3,2为5x5,类推
    :return:
    """
    print('medianblur processing...')
    dst = img.copy()
    black_pixels = []  # 图中像素值为0的点
    img = dst.copy()
    for r in range(0, img.shape[0]):
        for c in range(0, img.shape[1]):
            if img[r, c] == 0:
                black_pixels.append([r, c])
    black_pixels_temp = black_pixels
    while True:
        print(f'black_pixel:{len(black_pixels)}')
        for bp in black_pixels:
            pixel = []
            for i in range(max(0, bp[0]-kernel), min(img.shape[0], bp[0]+kernel+1)):
                for j in range(max(0, bp[1]-kernel), min(img.shape[1], bp[1]+kernel+1)):
                    if img[i, j] != 0:
                        pixel.append(img[i, j])  # 添加周围的非零像素点
            if len(pixel) == 0:
                continue
            black_pixels_temp.remove(bp)
            medium = np.array(pixel)
            dst[bp[0], bp[1]] = np.median(medium)
        black_pixels = black_pixels_temp
        img = dst
        if len(black_pixels_temp) == 0:
            return dst


def joint_filter(img, guideimg, kernel=5, sigma_space=1, sigma_color=1):
    print('jointfilter processing...')
    dst = img.copy()
    black_pixels = []  # 图中像素值为0的点
    img = dst.copy()
    for r in range(0, img.shape[0]):
        for c in range(0, img.shape[1]):
            if img[r, c] == 0:
                black_pixels.append([r, c])
    black_pixels_temp = black_pixels
    while True:
        print(f'black_pixel:{len(black_pixels)}')
        for bp in black_pixels:
            sum_weight = 0  # 权重和
            sum_pixel = 0  # 像素和
            pixel_cnt = 0
            for i in range(max(0, bp[0]-kernel), min(img.shape[0], bp[0]+kernel+1)):
                for j in range(max(0, bp[1]-kernel), min(img.shape[1], bp[1]+kernel+1)):
                    if img[i, j] != 0:
                        pixel_cnt += 1
                        exp_space = math.exp((-(i - bp[0]) ** 2 - (j - bp[1]) ** 2) / sigma_space)  # 空间权重
                        dist = (np.float32(guideimg[i, j]) - np.float32(guideimg[bp[0], bp[1]])) / 16
                        exp_color = math.exp(-np.inner(dist, dist) / sigma_color)  # 颜色权重
                        weight = exp_space*exp_color
                        sum_weight += weight
                        sum_pixel += weight*img[i, j]
            if pixel_cnt == 0:
                continue
            black_pixels_temp.remove(bp)
            dst[bp[0], bp[1]] = np.int16(sum_pixel / sum_weight + 0.5)
        black_pixels = black_pixels_temp
        img = dst
        if len(black_pixels_temp) == 0:
            return dst
